//
// Created by sunmy on 2020/10/17.
//

#ifndef HW1_GMRES_H
#define HW1_GMRES_H

#include "Eigen/Core"
#include <iostream>

template <int n = 10, int m = 10, typename T = long double>
void gmres_impl(T max_residual_norm = 0.0001, int max_iter_time = 100)
{
    using namespace Eigen;
    using namespace std;
    typedef Matrix<T, n, n> MatrixND;
    typedef Matrix<T, m + 1, m> MatrixMM1;
    typedef Matrix<T, m + 1, m + 1> MatrixM1D;
    typedef Matrix<T, n, m> MatrixNM;
    typedef Matrix<T, n, 1> VectorND;
    typedef Matrix<T, m, 1> VectorMD;
    typedef Matrix<T, m + 1, 1> VectorM1D;

    MatrixND A = MatrixND();

    for (auto i = 0; i < n; ++i)
    {
        for (auto j = 0; j < n; ++j)
        {
            A(i, j) = (T)1.0 / (i + j + 1);
        }
    }

    VectorND b = A * VectorND::Ones(n, 1);

    //    cout << "Now the matrix A is:" << endl << A << endl;
    //    cout << "Now the vector b is:" << endl << b << endl;

    // initial value
    VectorND x = VectorND::Zero();
    VectorND r = b - A * x;
    MatrixNM V = MatrixNM::Zero();
    MatrixMM1 H = MatrixMM1::Zero();

    // global iteration
    auto iter_time = 0;
    for (; iter_time < max_iter_time; ++iter_time)
    {
        V.col(0) = r / r.norm();
        for (auto j = 0; j < m; ++j)
        {
            for (auto i = 0; i <= j; ++i)
            {
                H(i, j) = (A * V.col(j)).dot(V.col(i));
            }
            VectorND tmp = A * V.col(j);
            for (auto i = 0; i <= j; ++i)
            {
                tmp -= H(i, j) * V.col(i);
            }
            H(j + 1, j) = tmp.norm();
            if (j < m - 1)
                V.col(j + 1) = tmp / tmp.norm();
        }

        // find y to minimize |beta e1-H y|
        VectorM1D betae1 = VectorM1D::Zero();
        betae1(0) = r.norm();

        // we need to perform rotations on betae1 and H to make H upper triangular
        for (auto j = 0; j < m; ++j)
        {
            MatrixM1D rotate = MatrixM1D::Identity();
            T c = H(j, j) / sqrt(H(j, j) * H(j, j) + H(j + 1, j) * H(j + 1, j));
            T s = -H(j + 1, j) / sqrt(H(j, j) * H(j, j) + H(j + 1, j) * H(j + 1, j));
            rotate(j, j) = c;
            rotate(j + 1, j + 1) = c;
            rotate(j, j + 1) = -s;
            rotate(j + 1, j) = s;

            H = rotate * H;
            betae1 = rotate * betae1;
        }

        // solve y from betae1 and H
        VectorMD y = VectorMD::Zero();
        for (auto i = m - 1; i >= 0; --i)
        {
            T sum = betae1(i);
            for (auto j = i + 1; j < m; ++j)
                sum -= H(i, j) * y(j);
            y(i) = sum / H(i, i);
        }

        // resolve x and r
        x += V * y;
        r = b - A * x;

        // theoritically, ||r|| should equal to betae1(m), we can save some computations
        if (betae1(m) < max_residual_norm)
            break;
    }

    cout.precision(12);
    cout << "When n = " << n << ", the iteration time is " << iter_time << endl
         << "with residual norm " << r.norm() << endl
         << "and the solution is:" << endl
         << x << endl;
}

void gmres()
{
    long double residual = 0.000000000000000001;
    gmres_impl<10>(residual);
    gmres_impl<11>(residual);
    gmres_impl<12>(residual);
    gmres_impl<13>(residual);
    gmres_impl<14>(residual);
    gmres_impl<15>(residual);
}

#endif //HW1_GMRES_H
